{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import tokenizers\n",
    "import string\n",
    "import torch\n",
    "import transformers\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from tqdm import tqdm\n",
    "import re"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 初期値"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "MAX_LEN = 192\n",
    "TRAIN_BATCH_SIZE = 32\n",
    "VALID_BATCH_SIZE = 8\n",
    "EPOCHS = 5\n",
    "ROBERTA_PATH = \"/home/tidal/ML_Data/Tweet_Sentiment_Extraction/_env_roberta-inference-5-folds/roberta-base\"\n",
    "#tokenizer:文章を最小単位であるtoken(字句)に分ける解析器\n",
    "TOKENIZER = tokenizers.ByteLevelBPETokenizer(\n",
    "    vocab_file=f\"{ROBERTA_PATH}/vocab.json\", \n",
    "    merges_file=f\"{ROBERTA_PATH}/merges.txt\", \n",
    "    lowercase=True,\n",
    "    add_prefix_space=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TweetModel(transformers.BertPreTrainedModel):\n",
    "    def __init__(self, conf):\n",
    "        super(TweetModel, self).__init__(conf)\n",
    "        self.roberta = transformers.RobertaModel.from_pretrained(ROBERTA_PATH, config=conf)\n",
    "        self.drop_out = nn.Dropout(0.1)\n",
    "        self.l0 = nn.Linear(768 * 2, 2)\n",
    "        torch.nn.init.normal_(self.l0.weight, std=0.02) #重みの初期化\n",
    "    \n",
    "    def forward(self, ids, mask, token_type_ids):\n",
    "        _, _, out = self.roberta(\n",
    "            ids,\n",
    "            attention_mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "\n",
    "        out = torch.cat((out[-1], out[-2]), dim=-1)\n",
    "        out = self.drop_out(out)\n",
    "        logits = self.l0(out)\n",
    "\n",
    "        start_logits, end_logits = logits.split(1, dim=-1)\n",
    "\n",
    "        start_logits = start_logits.squeeze(-1)\n",
    "        end_logits = end_logits.squeeze(-1)\n",
    "\n",
    "        return start_logits, end_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_data(tweet, selected_text, sentiment, tokenizer, max_len):\n",
    "    tweet = \" \" + \" \".join(str(tweet).split())\n",
    "    selected_text = \" \" + \" \".join(str(selected_text).split())\n",
    "\n",
    "    len_st = len(selected_text) - 1\n",
    "    idx0 = None\n",
    "    idx1 = None\n",
    "\n",
    "    for ind in (i for i, e in enumerate(tweet) if e == selected_text[1]):\n",
    "        if \" \" + tweet[ind: ind+len_st] == selected_text:\n",
    "            idx0 = ind\n",
    "            idx1 = ind + len_st - 1\n",
    "            break\n",
    "\n",
    "    char_targets = [0] * len(tweet)\n",
    "    if idx0 != None and idx1 != None:\n",
    "        for ct in range(idx0, idx1 + 1):\n",
    "            char_targets[ct] = 1\n",
    "    \n",
    "    tok_tweet = tokenizer.encode(tweet)\n",
    "    input_ids_orig = tok_tweet.ids\n",
    "    tweet_offsets = tok_tweet.offsets\n",
    "    \n",
    "    target_idx = []\n",
    "    for j, (offset1, offset2) in enumerate(tweet_offsets):\n",
    "        if sum(char_targets[offset1: offset2]) > 0:\n",
    "            target_idx.append(j)\n",
    "    \n",
    "    targets_start = target_idx[0]\n",
    "    targets_end = target_idx[-1]\n",
    "\n",
    "    sentiment_id = {\n",
    "        'positive': 1313,\n",
    "        'negative': 2430,\n",
    "        'neutral': 7974\n",
    "    }\n",
    "    \n",
    "    input_ids = [0] + [sentiment_id[sentiment]] + [2] + [2] + input_ids_orig + [2]\n",
    "    token_type_ids = [0, 0, 0, 0] + [0] * (len(input_ids_orig) + 1)\n",
    "    mask = [1] * len(token_type_ids)\n",
    "    tweet_offsets = [(0, 0)] * 4 + tweet_offsets + [(0, 0)]\n",
    "    targets_start += 4\n",
    "    targets_end += 4\n",
    "\n",
    "    padding_length = max_len - len(input_ids)\n",
    "    if padding_length > 0:\n",
    "        input_ids = input_ids + ([1] * padding_length)\n",
    "        mask = mask + ([0] * padding_length)\n",
    "        token_type_ids = token_type_ids + ([0] * padding_length)\n",
    "        tweet_offsets = tweet_offsets + ([(0, 0)] * padding_length)\n",
    "    \n",
    "    return {\n",
    "        'ids': input_ids,\n",
    "        'mask': mask,\n",
    "        'token_type_ids': token_type_ids,\n",
    "        'targets_start': targets_start,\n",
    "        'targets_end': targets_end,\n",
    "        'orig_tweet': tweet,\n",
    "        'orig_selected': selected_text,\n",
    "        'sentiment': sentiment,\n",
    "        'offsets': tweet_offsets\n",
    "    }\n",
    "\n",
    "\n",
    "class TweetDataset:\n",
    "    def __init__(self, tweet, sentiment, selected_text):\n",
    "        self.tweet = tweet\n",
    "        self.sentiment = sentiment\n",
    "        self.selected_text = selected_text\n",
    "        self.tokenizer = TOKENIZER\n",
    "        self.max_len = MAX_LEN\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.tweet)\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "        data = process_data(\n",
    "            self.tweet[item], \n",
    "            self.selected_text[item], \n",
    "            self.sentiment[item],\n",
    "            self.tokenizer,\n",
    "            self.max_len\n",
    "        )\n",
    "\n",
    "        return {\n",
    "            'ids': torch.tensor(data[\"ids\"], dtype=torch.long),\n",
    "            'mask': torch.tensor(data[\"mask\"], dtype=torch.long),\n",
    "            'token_type_ids': torch.tensor(data[\"token_type_ids\"], dtype=torch.long),\n",
    "            'targets_start': torch.tensor(data[\"targets_start\"], dtype=torch.long),\n",
    "            'targets_end': torch.tensor(data[\"targets_end\"], dtype=torch.long),\n",
    "            'orig_tweet': data[\"orig_tweet\"],\n",
    "            'orig_selected': data[\"orig_selected\"],\n",
    "            'sentiment': data[\"sentiment\"],\n",
    "            'offsets': torch.tensor(data[\"offsets\"], dtype=torch.long)\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_jaccard_score(\n",
    "    original_tweet, \n",
    "    target_string, \n",
    "    sentiment_val, \n",
    "    idx_start, \n",
    "    idx_end, \n",
    "    offsets,\n",
    "    verbose=False):\n",
    "    \n",
    "    if idx_end < idx_start:\n",
    "        idx_end = idx_start\n",
    "    \n",
    "    filtered_output  = \"\"\n",
    "    for ix in range(idx_start, idx_end + 1):\n",
    "        filtered_output += original_tweet[offsets[ix][0]: offsets[ix][1]]\n",
    "        if (ix+1) < len(offsets) and offsets[ix][1] < offsets[ix+1][0]:\n",
    "            filtered_output += \" \"\n",
    "\n",
    "    if sentiment_val == \"neutral\" or len(original_tweet.split()) < 2:\n",
    "        filtered_output = original_tweet\n",
    "\n",
    "    if sentiment_val != \"neutral\" and verbose == True:\n",
    "        if filtered_output.strip().lower() != target_string.strip().lower():\n",
    "            print(\"********************************\")\n",
    "            print(f\"Output= {filtered_output.strip()}\")\n",
    "            print(f\"Target= {target_string.strip()}\")\n",
    "            print(f\"Tweet= {original_tweet.strip()}\")\n",
    "            print(\"********************************\")\n",
    "\n",
    "    jac = 0\n",
    "    return jac, filtered_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = pd.read_csv(\"/home/tidal/ML_Data/Tweet_Sentiment_Extraction/test.csv\")\n",
    "df_test.loc[:, \"selected_text\"] = df_test.text.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "model_config = transformers.RobertaConfig.from_pretrained(ROBERTA_PATH)\n",
    "model_config.output_hidden_states = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TweetModel(\n",
       "  (roberta): RobertaModel(\n",
       "    (embeddings): RobertaEmbeddings(\n",
       "      (word_embeddings): Embedding(50265, 768, padding_idx=1)\n",
       "      (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "      (token_type_embeddings): Embedding(1, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (drop_out): Dropout(p=0.1, inplace=False)\n",
       "  (l0): Linear(in_features=1536, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model1 = TweetModel(conf=model_config)\n",
    "model1.to(device)\n",
    "model1.load_state_dict(torch.load(\"../input/tweet-roberta/model_0.bin\"))\n",
    "model1.eval()\n",
    "\n",
    "model2 = TweetModel(conf=model_config)\n",
    "model2.to(device)\n",
    "model2.load_state_dict(torch.load(\"../input/tweet-roberta/model_1.bin\"))\n",
    "model2.eval()\n",
    "\n",
    "model3 = TweetModel(conf=model_config)\n",
    "model3.to(device)\n",
    "model3.load_state_dict(torch.load(\"../input/tweet-roberta/model_2.bin\"))\n",
    "model3.eval()\n",
    "\n",
    "model4 = TweetModel(conf=model_config)\n",
    "model4.to(device)\n",
    "model4.load_state_dict(torch.load(\"../input/tweet-roberta/model_3.bin\"))\n",
    "model4.eval()\n",
    "\n",
    "model5 = TweetModel(conf=model_config)\n",
    "model5.to(device)\n",
    "model5.load_state_dict(torch.load(\"../input/tweet-roberta/model_4.bin\"))\n",
    "model5.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_output = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 442/442 [01:52<00:00,  3.92it/s]\n"
     ]
    }
   ],
   "source": [
    "test_dataset = TweetDataset(\n",
    "        tweet=df_test.text.values,\n",
    "        sentiment=df_test.sentiment.values,\n",
    "        selected_text=df_test.selected_text.values\n",
    "    )\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    test_dataset,\n",
    "    shuffle=False,\n",
    "    batch_size=VALID_BATCH_SIZE,\n",
    "    num_workers=1\n",
    ")\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    tk0 = tqdm(data_loader, total=len(data_loader))\n",
    "    for bi, d in enumerate(tk0):\n",
    "        ids = d[\"ids\"]\n",
    "        token_type_ids = d[\"token_type_ids\"]\n",
    "        mask = d[\"mask\"]\n",
    "        sentiment = d[\"sentiment\"]\n",
    "        orig_selected = d[\"orig_selected\"]\n",
    "        orig_tweet = d[\"orig_tweet\"]\n",
    "        targets_start = d[\"targets_start\"]\n",
    "        targets_end = d[\"targets_end\"]\n",
    "        offsets = d[\"offsets\"].numpy()\n",
    "\n",
    "        ids = ids.to(device, dtype=torch.long)\n",
    "        token_type_ids = token_type_ids.to(device, dtype=torch.long)\n",
    "        mask = mask.to(device, dtype=torch.long)\n",
    "        targets_start = targets_start.to(device, dtype=torch.long)\n",
    "        targets_end = targets_end.to(device, dtype=torch.long)\n",
    "\n",
    "        outputs_start1, outputs_end1 = model1(\n",
    "            ids=ids,\n",
    "            mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        \n",
    "        outputs_start2, outputs_end2 = model2(\n",
    "            ids=ids,\n",
    "            mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        \n",
    "        outputs_start3, outputs_end3 = model3(\n",
    "            ids=ids,\n",
    "            mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        \n",
    "        outputs_start4, outputs_end4 = model4(\n",
    "            ids=ids,\n",
    "            mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        \n",
    "        outputs_start5, outputs_end5 = model5(\n",
    "            ids=ids,\n",
    "            mask=mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )\n",
    "        outputs_start = (outputs_start1 + outputs_start2 + outputs_start3 + outputs_start4 + outputs_start5) / 5\n",
    "        outputs_end = (outputs_end1 + outputs_end2 + outputs_end3 + outputs_end4 + outputs_end5) / 5\n",
    "        \n",
    "        outputs_start = torch.softmax(outputs_start, dim=1).cpu().detach().numpy()\n",
    "        outputs_end = torch.softmax(outputs_end, dim=1).cpu().detach().numpy()\n",
    "        jaccard_scores = []\n",
    "        for px, tweet in enumerate(orig_tweet):\n",
    "            selected_tweet = orig_selected[px]\n",
    "            tweet_sentiment = sentiment[px]\n",
    "            _, output_sentence = calculate_jaccard_score(\n",
    "                original_tweet=tweet,\n",
    "                target_string=selected_tweet,\n",
    "                sentiment_val=tweet_sentiment,\n",
    "                idx_start=np.argmax(outputs_start[px, :]),\n",
    "                idx_end=np.argmax(outputs_end[px, :]),\n",
    "                offsets=offsets[px]\n",
    "            )\n",
    "            final_output.append(output_sentence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def post_process(selected):\n",
    "    return \" \".join(set(selected.lower().split()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = pd.read_csv(\"../input/tweet-sentiment-extraction/sample_submission.csv\")\n",
    "sample.loc[:, 'selected_text'] = final_output\n",
    "sample.selected_text = sample.selected_text.map(post_process)\n",
    "sample.to_csv(\"submission.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>textID</th>\n",
       "      <th>selected_text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>11aa4945ff</td>\n",
       "      <td>wish i</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>fd1db57dc0</td>\n",
       "      <td>done.haha. i'm</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2524332d66</td>\n",
       "      <td>concerned that family i'm for</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0fb19285b2</td>\n",
       "      <td>worry. no to need</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>e6c9e5e3ab</td>\n",
       "      <td>26th february</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       textID                  selected_text\n",
       "0  11aa4945ff                         wish i\n",
       "1  fd1db57dc0                 done.haha. i'm\n",
       "2  2524332d66  concerned that family i'm for\n",
       "3  0fb19285b2              worry. no to need\n",
       "4  e6c9e5e3ab                  26th february"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
